// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/plan/plan_builder.h>
#include <pollux/runner/multi_fragment_plan.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    /// Builder for distributed plans. Adds a shuffle() and related
    /// methods for building PartitionedOutput-Exchange pairs between
    /// fragments. Not thread safe.
    class DistributedPlanBuilder : public PlanBuilder {
    public:
        /// Constructs a top level DistributedPlanBuilder.
        DistributedPlanBuilder(
            const runner::MultiFragmentPlan::Options &options,
            std::shared_ptr<core::PlanNodeIdGenerator> planNodeIdGenerator,
            memory::MemoryPool *pool = nullptr);

        /// Constructs a child builder. Used for branching plans, e.g. the subplan for
        /// a join build side.
        DistributedPlanBuilder(DistributedPlanBuilder &root);

        ~DistributedPlanBuilder() override;

        /// Returns the planned fragments. The builder will be empty after this. This
        /// is only called on the root builder.
        std::vector<runner::ExecutableFragment> fragments();

        PlanBuilder &shufflePartitioned(
            const std::vector<std::string> &keys,
            int numPartitions,
            bool replicateNullsAndAny,
            const std::vector<std::string> &outputLayout = {}) override;

        core::PlanNodePtr shufflePartitionedResult(
            const std::vector<std::string> &keys,
            int numPartitions,
            bool replicateNullsAndAny,
            const std::vector<std::string> &outputLayout = {}) override;

        core::PlanNodePtr shuffleBroadcastResult() override;

    private:
        void newFragment();

        void gatherScans(const core::PlanNodePtr &plan);

        const runner::MultiFragmentPlan::Options &options_;
        DistributedPlanBuilder *const root_;

        // Stack of outstanding builders. The last element is the immediately
        // enclosing one. When returning an ExchangeNode from returnShuffle, the stack
        // is used to establish the linkage between the fragment of the returning
        // builder and the fragment current in the calling builder. Only filled in the
        // root builder.
        std::vector<DistributedPlanBuilder *> stack_;

        // Fragment counter. Only used in root builder.
        int32_t fragmentCounter_{0};

        // The fragment being built. Will be moved to the root builder's 'fragments_'
        // when complete.
        std::unique_ptr<runner::ExecutableFragment> current_;

        // The fragments gathered under this builder. Moved to the root builder when
        // returning the subplan.
        std::vector<runner::ExecutableFragment> fragments_;
    };
} // namespace kumo::pollux::exec::test
